#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
批量YOLO格式转VOC格式转换脚本
将251003目录下的所有YOLO格式数据集转换为VOC格式
"""

import os
import sys
from pathlib import Path
from cov_yolo_voc import YOLOToVOCConverter


def batch_convert_yolo_to_voc():
    """批量转换YOLO格式数据集为VOC格式"""

    # 源目录和目标目录
    source_dir = Path(r"F:\GitProject\yolov7-pytorch\VOCdevkit\251003")
    target_dir = Path(r"F:\GitProject\yolov7-pytorch\VOCdevkit\251003voc")

    # 确保目标目录存在
    target_dir.mkdir(parents=True, exist_ok=True)

    # 获取所有YOLO格式的文件夹
    yolo_folders = []
    for item in source_dir.iterdir():
        if item.is_dir() and not item.name.endswith("_voc"):
            # 检查是否包含YOLO格式文件
            if (item / "obj.data").exists() and (item / "obj.names").exists():
                yolo_folders.append(item)

    print(f"找到 {len(yolo_folders)} 个YOLO格式数据集:")
    for folder in yolo_folders:
        print(f"  - {folder.name}")

    print("\n开始批量转换...")

    # 转换每个数据集
    for i, yolo_folder in enumerate(yolo_folders, 1):
        print(f"\n[{i}/{len(yolo_folders)}] 转换数据集: {yolo_folder.name}")

        # 设置输入和输出路径
        yolo_dir = yolo_folder
        output_dir = target_dir / f"{yolo_folder.name}_voc"
        class_names_file = yolo_folder / "obj.names"

        print(f"  输入目录: {yolo_dir}")
        print(f"  输出目录: {output_dir}")
        print(f"  类别文件: {class_names_file}")

        try:
            # 创建转换器
            converter = YOLOToVOCConverter(
                yolo_dir=str(yolo_dir),
                output_dir=str(output_dir),
                class_names_file=str(class_names_file),
            )

            # 执行转换
            converter.convert_dataset()

            print(f"  [OK] {yolo_folder.name} 转换完成")

        except Exception as e:
            print(f"  [ERROR] {yolo_folder.name} 转换失败: {e}")
            continue

    print(f"\n[SUCCESS] 批量转换完成！")
    print(f"所有VOC格式数据集保存在: {target_dir}")

    # 显示转换结果
    print(f"\n转换结果:")
    for item in target_dir.iterdir():
        if item.is_dir() and item.name.endswith("_voc"):
            annotations_count = (
                len(list((item / "Annotations").glob("*.xml")))
                if (item / "Annotations").exists()
                else 0
            )
            images_count = (
                len(list((item / "JPEGImages").glob("*.jpg")))
                if (item / "JPEGImages").exists()
                else 0
            )
            print(
                f"  - {item.name}: {images_count} 张图片, {annotations_count} 个标注文件"
            )


if __name__ == "__main__":
    batch_convert_yolo_to_voc()
